import numpy as np
from functools import reduce
import cv2
import os
import shutil

from argparse import ArgumentParser

def fine_jpg_in_Dir_recursively( myDir ):
  files = os.listdir( myDir )
  jpgs = []
  for f in files:
    f = os.path.join( myDir , f )
    if not os.path.isdir( f ) and f[ -4: ] == ".jpg":
      jpgs.append( f )
    if os.path.isdir( f ):
      jpgs.extend( fine_jpg_in_Dir_recursively( f ) )
  return jpgs

def jpg_rename():
  srcDir = "/home/jh/working_data/idCard_v2"
  files = fine_jpg_in_Dir_recursively ( srcDir )

  negTarDir = "/home/jh/working_data/idData/neg"
  posTarDir = "/home/jh/working_data/idData/pos"

  for f in files:
    f_info = f.split( '/' )

    if f_info[-2] == '0':
      fname = "{:0>10d}".format( len( os.listdir( posTarDir ) ) ) + \
          "_pos_" + "base" + ".jpg"
      tarFile = os.path.join( posTarDir , fname )
      shutil.copyfile( f , tarFile )

    if 'PIC' in f_info[-1]:
      fname = "{:0>10d}".format( len( os.listdir( posTarDir ) ) ) + \
          "_pos_" + "fromICBC" + ".jpg"
      tarFile = os.path.join( posTarDir , fname )
      shutil.copyfile( f , tarFile )

    if f_info[-2] == '1':
      fname = "{:0>10d}".format( len( os.listdir( negTarDir ) ) ) + \
          "_neg_" + "base" + ".jpg"
      tarFile = os.path.join( negTarDir , fname )
      shutil.copyfile( f , tarFile )

    if 'ID' in f_info[-1]:
      fname = "{:0>10d}".format( len( os.listdir( negTarDir ) ) ) + \
          "_neg_" + "fromICBC" + ".jpg"
      tarFile = os.path.join( negTarDir , fname )
      shutil.copyfile( f , tarFile )

  """
  tarDir = "/home/jh/working_data/idData/neg"

  files = os.listdir( srcDir )

  for f in files:
    fname = "{:0>10d}".format( len( os.listdir(tarDir) ) ) + \
        "_neg_" + "monitor" + ".jpg"
    tarFile = os.path.join( tarDir , fname )
    srcFile = os.path.join( srcDir , f )
    shutil.copyfile( srcFile , tarFile )
  """

def change_name():
  data_dir = "/home/jh/working_data/FakeImagesByPicturingFromMonitor" 
  fileList = os.listdir( data_dir )

  for f in fileList:
    srcFile = os.path.join( data_dir , f )
    tarFile = os.path.join( data_dir , "Fake_" + f )
    os.rename( srcFile , tarFile )

def parse_args():
  parser = ArgumentParser( description = "first time import such a module" )

  parser.add_argument( '-i', '--info' , 
      dest = "info_path" , 
      default = './id_data.txt' , help = "txt file from mtcnn" )
  parser.add_argument( '--data' , dest = "data_path" , 
      default = "/Users/pitaloveu/working_data/idCard" )
  parser.add_argument( '--write_path' , dest = "write_path" , 
      default = "/Users/pitaloveu/working_data/idCard/extract" )

  return parser.parse_args()

def run_face_extraction():
  """
  1. read all the infomation from the txt files
    including 0 - many bounding boxes coordinates in the file
    coordinates separated by "|"

    boxes coordinates seems in a form : 
    N * x1 , N * y1 , N * x2 , N * y2 , N * score

  2. the whole idCard data will mainly contain two parts:
    (1). directories pos/ contains all positive images, including the ones from following sources:
  a. ICBC
  b. from Li lei
  c. base data from Zhanglei
  (2). directories neg/ contains all the negative imgaes
  """

  txt_path   = parse_args().info_path
  data_path  = parse_args().data_path
  write_path = parse_args().write_path

  print( "this is the txt file %s " % txt_path )

  if not os.path.isfile( txt_path ):
    raise ValueError( "buddy, input wrong" )

  with open( txt_path , 'r' ) as fo:
    lines = fo.readlines()
  
  # filter all the empty strings
  lines = [ line for line in lines if \
      line.split( '|' )[-1].strip() != "" ]

  def do_something( line ):
    """
    this function is used to do some clipping or showing
    """
    line_split = list( map ( lambda s: s.strip() , line.split( '|' ) ))
    bboxes = list( map( lambda s: float( s ) , line_split[1].split( ' ' ) ) )
    num_bboxes = int( len( bboxes ) / 5 )
    img_name = os.path.basename( line_split[0] )

    if 'neg' == line_split[0].split('/')[-2]:
      abs_img_name = os.path.join( data_path , 'neg' , img_name )
    else:
      abs_img_name = os.path.join( data_path , 'pos' , img_name )

    img = cv2.imread( abs_img_name )

    height = img.shape[0]
    width  = img.shape[1]

    # drawing bboxes in the image to testify 
    # the correctness of mtcnn face detection

    max_area = 0
    for index in range( num_bboxes ):
      bb = [ bboxes[index] , bboxes[ num_bboxes + index ] , \
          bboxes[ 2 * num_bboxes + index ] ,\
          bboxes[ 3 * num_bboxes + index ] ]
      bb = list( map( lambda x :int( x) , bb ) )

      area = ( bb[2] - bb[0] ) * ( bb[3] - bb[1] )
      if area > max_area:
        max_area = area
        max_bb = bb
      
    #cv2.rectangle( img , ( max_bb[0] , max_bb[1] ) , \
    #    ( max_bb[2] , max_bb[3] ) , (0,255,0) , 2 )

    max_bb[1] = max( 0 , max_bb[1] )
    max_bb[3] = min( max_bb[3] , height-1 )
    max_bb[0] = max( 0 , max_bb[0] )
    max_bb[2] = min( max_bb[2] , width-1 )

    # add some extending pixels
    tmpH = max_bb[3] - max_bb[1]
    tmpW = max_bb[2] - max_bb[0]

    max_bb[0] = int( max_bb[0] - 0.5 * tmpW )
    max_bb[2] = int( max_bb[2] + 0.5 * tmpW )
    max_bb[1] = int( max_bb[1] - 0.5 * tmpH )
    max_bb[3] = int( max_bb[3] + 0.5 * tmpH )

    
    max_bb[1] = max( 0 , max_bb[1] )
    max_bb[3] = min( max_bb[3] , height-1 )
    max_bb[0] = max( 0 , max_bb[0] )
    max_bb[2] = min( max_bb[2] , width-1 )

    img_output = img[ max_bb[1] : max_bb[3] , \
        max_bb[0] : max_bb[2] ]

    if 'neg' == line_split[0].split('/')[-2]:
      output_path = os.path.join( write_path , "neg" , img_name )
    else:
      output_path = os.path.join( write_path , "pos" , img_name )

    cv2.imwrite( output_path , img_output )
     
  for line in lines :
    do_something( line )

def run_face_extraction_fakeImg():
  """
  1. read all the infomation from the txt files
    including 0 - many bounding boxes coordinates in the file
    coordinates separated by "|"

    boxes coordinates seems in a form : 
    N * x1 , N * y1 , N * x2 , N * y2 , N * score

  2. all the data in the FakeImage... Folder all is fake

  """

  txt_path   = parse_args().info_path
  data_path  = parse_args().data_path
  write_path = parse_args().write_path

  print( "this is the txt file %s " % txt_path )

  if not os.path.isfile( txt_path ):
    raise ValueError( "buddy, input wrong" )

  with open( txt_path , 'r' ) as fo:
    lines = fo.readlines()
  
  # filter all the empty strings
  lines = [ line for line in lines if \
      line.split( '|' )[-1].strip() != "" ]

  def do_something( line ):
    """
    this function is used to do some clipping or showing
    """
    line_split = list( map ( lambda s: s.strip() , line.split( '|' ) ))
    bboxes = list( map( lambda s: float( s ) , line_split[1].split( ' ' ) ) )
    num_bboxes = int( len( bboxes ) / 5 )
    img_name = os.path.basename( line_split[0] )

    abs_img_name = os.path.join( data_path , img_name )

    img = cv2.imread( abs_img_name )

    height = img.shape[0]
    width  = img.shape[1]

    # drawing bboxes in the image to testify 
    # the correctness of mtcnn face detection

    max_area = 0
    for index in range( num_bboxes ):
      bb = [ bboxes[index] , bboxes[ num_bboxes + index ] , \
          bboxes[ 2 * num_bboxes + index ] ,\
          bboxes[ 3 * num_bboxes + index ] ]
      bb = list( map( lambda x :int( x) , bb ) )

      area = ( bb[2] - bb[0] ) * ( bb[3] - bb[1] )
      if area > max_area:
        max_area = area
        max_bb = bb
      
    #cv2.rectangle( img , ( max_bb[0] , max_bb[1] ) , \
    #    ( max_bb[2] , max_bb[3] ) , (0,255,0) , 2 )

    max_bb[1] = max( 0 , max_bb[1] )
    max_bb[3] = min( max_bb[3] , height-1 )
    max_bb[0] = max( 0 , max_bb[0] )
    max_bb[2] = min( max_bb[2] , width-1 )

    # add some extending pixels
    tmpH = max_bb[3] - max_bb[1]
    tmpW = max_bb[2] - max_bb[0]

    max_bb[0] = int( max_bb[0] - 0.5 * tmpW )
    max_bb[2] = int( max_bb[2] + 0.5 * tmpW )
    max_bb[1] = int( max_bb[1] - 0.5 * tmpH )
    max_bb[3] = int( max_bb[3] + 0.5 * tmpH )

    
    max_bb[1] = max( 0 , max_bb[1] )
    max_bb[3] = min( max_bb[3] , height-1 )
    max_bb[0] = max( 0 , max_bb[0] )
    max_bb[2] = min( max_bb[2] , width-1 )


    img_output = img[ max_bb[1] : max_bb[3] , \
        max_bb[0] : max_bb[2] ]


    output_path = os.path.join( write_path , img_name )

    cv2.imwrite( output_path , img_output )
     
  for line in lines :
    do_something( line )

def run():
  """
  the workflow will be:
  1. load the information of all the labeled data

  2. for every single element in the query sequence, find the one
     in the labeled data

  3. normalize the landmark first, then compute the pixel error

  4. Loop all the qualified query one in, and average them
  """
  image_path , landmarks, gender, smile, glasses, pose = \
      fetchData.load_path( data_path , if_train = False )

  imdb = list( zip( image_path , landmarks ) )

  with open( 'mtcnn_test_results.txt' , 'r' ) as f:
    lines = f.readlines()

  lines = [ line for line in lines if \
      line.split( '|' )[-1].strip() != "" and \
      len( line.split( '|' )[-1].strip().split( ' ' ) ) == 10 ]

  errors = []
  for line in lines:
    line_split = list( map ( lambda s: s.strip() , line.split( '|' ) ))
    _key = line_split[0].split( '/' )[-1]
    predict_landmarks = list( map( lambda s: float( s ) , line_split[-1].split( ' ' ) ) )

    imdb_info = [ s for s in imdb if s[0].find(_key) != -1 ]
    assert len( imdb_info ) == 1

    img = cv2.imread( imdb_info[0][0] )
    labeled_landmarks = imdb_info[0][1]
    height = img.shape[0]
    width  = img.shape[1]

    # transfer the landmark labeled to ( 96, 112 ) format
    labeled_landmarks[0:5] = labeled_landmarks[0:5] * 96. / width
    labeled_landmarks[5:10] = labeled_landmarks[5:10] * 112. / height

    #error in a form [ x1, x2 , x3, x4, x5, y1, y2, y3, y4, y5 ]
    error_landmarks = np.array( predict_landmarks ) - labeled_landmarks
    error = np.sqrt( np.square( error_landmarks[0:5]) + \
        np.square( error_landmarks[5:10]) )

    errors.append( error )

  errors = np.array( errors )
  errors_mean = np.mean( errors , 0 )
  print( errors_mean )

if __name__ == "__main__":
  #run_face_extraction_fakeImg()
  run_face_extraction()

  #jpg_rename()

